"""
WebSocket实时协作服务
提供实时通信和协作功能
"""

import json
import asyncio
from typing import Dict, Set, Any, Optional
from fastapi import WebSocket, WebSocketDisconnect
from src.utils.logging import get_logger

logger = get_logger(__name__)

class ConnectionManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        # 存储活跃连接：{session_id: {user_id: websocket}}
        self.active_connections: Dict[str, Dict[str, WebSocket]] = {}
        # 存储用户信息：{user_id: user_info}
        self.user_info: Dict[str, Dict[str, Any]] = {}
        # 存储会话研究状态：{session_id: research_state}
        self.session_research: Dict[str, Dict[str, Any]] = {}
    
    async def connect(self, websocket: WebSocket, session_id: str, user_id: str, user_info: Optional[Dict[str, Any]] = None):
        """
        建立WebSocket连接
        
        Args:
            websocket: WebSocket连接对象
            session_id: 会话ID
            user_id: 用户ID
            user_info: 用户信息
        """
        await websocket.accept()
        
        # 初始化会话连接字典（如果不存在）
        if session_id not in self.active_connections:
            self.active_connections[session_id] = {}
        
        # 存储连接
        self.active_connections[session_id][user_id] = websocket
        self.user_info[user_id] = user_info or {"user_id": user_id}
        
        logger.info(f"用户 {user_id} 连接到会话 {session_id}")
        
        # 发送欢迎消息
        await self.send_personal_message(
            {"type": "connected", "message": "成功连接到协作会话"}, 
            user_id, 
            session_id
        )
        
        # 通知其他用户有新用户加入
        await self.broadcast_to_session(
            {
                "type": "user_joined",
                "user_id": user_id,
                "user_info": self.user_info[user_id],
                "message": f"用户 {user_id} 加入了会话"
            },
            session_id,
            exclude_user=user_id
        )
        
        # 如果会话中有正在进行的研究，发送当前研究状态
        if session_id in self.session_research:
            await self.send_personal_message(
                {
                    "type": "research_state",
                    "state": self.session_research[session_id]
                },
                user_id,
                session_id
            )
    
    def disconnect(self, session_id: str, user_id: str):
        """
        断开WebSocket连接
        
        Args:
            session_id: 会话ID
            user_id: 用户ID
        """
        if session_id in self.active_connections:
            if user_id in self.active_connections[session_id]:
                del self.active_connections[session_id][user_id]
                logger.info(f"用户 {user_id} 从会话 {session_id} 断开连接")
                
                # 如果会话中没有其他用户，清理会话
                if not self.active_connections[session_id]:
                    del self.active_connections[session_id]
                    # 清理会话研究状态
                    if session_id in self.session_research:
                        del self.session_research[session_id]
                    logger.info(f"会话 {session_id} 已关闭")
            
            # 通知其他用户有用户离开
            if session_id in self.active_connections:
                asyncio.create_task(
                    self.broadcast_to_session(
                        {
                            "type": "user_left",
                            "user_id": user_id,
                            "message": f"用户 {user_id} 离开了会话"
                        },
                        session_id,
                        exclude_user=user_id
                    )
                )
    
    async def send_personal_message(self, message: Dict[str, Any], user_id: str, session_id: str):
        """
        发送个人消息
        
        Args:
            message: 消息内容
            user_id: 用户ID
            session_id: 会话ID
        """
        try:
            if (session_id in self.active_connections and 
                user_id in self.active_connections[session_id]):
                websocket = self.active_connections[session_id][user_id]
                await websocket.send_text(json.dumps(message, ensure_ascii=False))
        except Exception as e:
            logger.error(f"发送个人消息失败: {e}")
            self.disconnect(session_id, user_id)
    
    async def broadcast_to_session(self, message: Dict[str, Any], session_id: str, exclude_user: Optional[str] = None):
        """
        向会话中的所有用户广播消息
        
        Args:
            message: 消息内容
            session_id: 会话ID
            exclude_user: 排除的用户ID（可选）
        """
        if session_id not in self.active_connections:
            return
        
        disconnected_users = []
        
        for user_id, websocket in self.active_connections[session_id].items():
            # 跳过排除的用户
            if user_id == exclude_user:
                continue
                
            try:
                await websocket.send_text(json.dumps(message, ensure_ascii=False))
            except Exception as e:
                logger.error(f"向用户 {user_id} 发送消息失败: {e}")
                disconnected_users.append(user_id)
        
        # 清理断开连接的用户
        for user_id in disconnected_users:
            self.disconnect(session_id, user_id)
    
    def update_research_state(self, session_id: str, state: Dict[str, Any]):
        """
        更新会话研究状态
        
        Args:
            session_id: 会话ID
            state: 研究状态
        """
        self.session_research[session_id] = state
    
    def get_session_users(self, session_id: str) -> Dict[str, Dict[str, Any]]:
        """
        获取会话中的所有用户
        
        Args:
            session_id: 会话ID
            
        Returns:
            用户信息字典
        """
        if session_id not in self.active_connections:
            return {}
        
        users = {}
        for user_id in self.active_connections[session_id]:
            users[user_id] = self.user_info.get(user_id, {"user_id": user_id})
        
        return users

# 全局连接管理器实例
manager = ConnectionManager()

class WebSocketService:
    """WebSocket服务类"""
    
    def __init__(self):
        self.connection_manager = manager
        logger.info("WebSocket服务初始化完成")
    
    async def handle_websocket_connection(self, websocket: WebSocket, session_id: str, user_id: str, user_info: Optional[Dict[str, Any]] = None):
        """
        处理WebSocket连接
        
        Args:
            websocket: WebSocket连接对象
            session_id: 会话ID
            user_id: 用户ID
            user_info: 用户信息
        """
        # 建立连接
        await self.connection_manager.connect(websocket, session_id, user_id, user_info)
        
        try:
            while True:
                # 接收客户端消息
                data = await websocket.receive_text()
                message = json.loads(data)
                
                # 处理不同类型的消息
                await self.process_message(message, session_id, user_id)
                
        except WebSocketDisconnect:
            logger.info(f"WebSocket连接断开: 用户 {user_id}, 会话 {session_id}")
        except Exception as e:
            logger.error(f"处理WebSocket消息时出错: {e}")
        finally:
            # 清理连接
            self.connection_manager.disconnect(session_id, user_id)
    
    async def process_message(self, message: Dict[str, Any], session_id: str, user_id: str):
        """
        处理接收到的消息
        
        Args:
            message: 消息内容
            session_id: 会话ID
            user_id: 用户ID
        """
        message_type = message.get("type", "unknown")
        
        # 添加发送者信息
        message["sender_id"] = user_id
        message["sender_info"] = self.connection_manager.user_info.get(user_id, {"user_id": user_id})
        
        if message_type == "chat_message":
            # 处理聊天消息
            await self.handle_chat_message(message, session_id, user_id)
        elif message_type == "research_update":
            # 处理研究更新
            await self.handle_research_update(message, session_id)
        elif message_type == "cursor_position":
            # 处理光标位置更新
            await self.handle_cursor_update(message, session_id, user_id)
        elif message_type == "research_start":
            # 处理研究开始请求
            await self.handle_research_start(message, session_id, user_id)
        elif message_type == "research_stop":
            # 处理研究停止请求
            await self.handle_research_stop(message, session_id, user_id)
        elif message_type == "research_query":
            # 处理研究查询
            await self.handle_research_query(message, session_id, user_id)
        else:
            # 广播未知类型消息
            await self.connection_manager.broadcast_to_session(
                message, 
                session_id, 
                exclude_user=user_id
            )
    
    async def handle_chat_message(self, message: Dict[str, Any], session_id: str, user_id: str):
        """
        处理聊天消息
        
        Args:
            message: 消息内容
            session_id: 会话ID
            user_id: 用户ID
        """
        # 广播聊天消息给会话中的所有其他用户
        await self.connection_manager.broadcast_to_session(
            {
                "type": "chat_message",
                "sender_id": user_id,
                "sender_info": message["sender_info"],
                "content": message.get("content", ""),
                "timestamp": message.get("timestamp")
            },
            session_id,
            exclude_user=user_id
        )
    
    async def handle_research_update(self, message: Dict[str, Any], session_id: str):
        """
        处理研究更新
        
        Args:
            message: 消息内容
            session_id: 会话ID
        """
        # 更新会话研究状态
        research_state = message.get("state", {})
        self.connection_manager.update_research_state(session_id, research_state)
        
        # 广播研究更新给会话中的所有用户
        await self.connection_manager.broadcast_to_session(
            {
                "type": "research_update",
                "sender_id": message["sender_id"],
                "sender_info": message["sender_info"],
                "question": message.get("question"),
                "answer": message.get("answer"),
                "status": message.get("status"),
                "state": research_state,
                "timestamp": message.get("timestamp")
            },
            session_id
        )
    
    async def handle_cursor_update(self, message: Dict[str, Any], session_id: str, user_id: str):
        """
        处理光标位置更新
        
        Args:
            message: 消息内容
            session_id: 会话ID
            user_id: 用户ID
        """
        # 广播光标位置更新给会话中的所有其他用户
        await self.connection_manager.broadcast_to_session(
            {
                "type": "cursor_position",
                "sender_id": user_id,
                "sender_info": message["sender_info"],
                "position": message.get("position"),
                "timestamp": message.get("timestamp")
            },
            session_id,
            exclude_user=user_id
        )
    
    async def handle_research_start(self, message: Dict[str, Any], session_id: str, user_id: str):
        """
        处理研究开始请求
        
        Args:
            message: 消息内容
            session_id: 会话ID
            user_id: 用户ID
        """
        # 广播研究开始消息
        await self.connection_manager.broadcast_to_session(
            {
                "type": "research_start",
                "sender_id": user_id,
                "sender_info": message["sender_info"],
                "question": message.get("question"),
                "timestamp": message.get("timestamp")
            },
            session_id
        )
    
    async def handle_research_stop(self, message: Dict[str, Any], session_id: str, user_id: str):
        """
        处理研究停止请求
        
        Args:
            message: 消息内容
            session_id: 会话ID
            user_id: 用户ID
        """
        # 广播研究停止消息
        await self.connection_manager.broadcast_to_session(
            {
                "type": "research_stop",
                "sender_id": user_id,
                "sender_info": message["sender_info"],
                "reason": message.get("reason"),
                "timestamp": message.get("timestamp")
            },
            session_id
        )
    
    async def handle_research_query(self, message: Dict[str, Any], session_id: str, user_id: str):
        """
        处理研究查询
        
        Args:
            message: 消息内容
            session_id: 会话ID
            user_id: 用户ID
        """
        # 广播研究查询消息
        await self.connection_manager.broadcast_to_session(
            {
                "type": "research_query",
                "sender_id": user_id,
                "sender_info": message["sender_info"],
                "query": message.get("query"),
                "timestamp": message.get("timestamp")
            },
            session_id,
            exclude_user=user_id
        )
    
    def get_session_users(self, session_id: str) -> Dict[str, Dict[str, Any]]:
        """
        获取会话中的所有用户
        
        Args:
            session_id: 会话ID
            
        Returns:
            用户信息字典
        """
        return self.connection_manager.get_session_users(session_id)
    
    def update_research_state(self, session_id: str, state: Dict[str, Any]):
        """
        更新会话研究状态
        
        Args:
            session_id: 会话ID
            state: 研究状态
        """
        self.connection_manager.update_research_state(session_id, state)

# 全局WebSocket服务实例
websocket_service = WebSocketService()